import asyncio
import time
from datetime import datetime

from langchain_core.messages import HumanMessage

from graph_agent.charts_agent.agent import create_charts_react_agent


async def main():
    start_time = time.perf_counter()
    agent = await create_charts_react_agent()
    # 示例查询
    queries = [
        # "统计今年的油单数量",
        "当前时间为2024-06-18，查询最近半年每个月的油单数量",
    ]

    for query in queries:
        state = {"messages": [HumanMessage(content=query)]}
        final_state = await agent.ainvoke(state)
        print(f"Query: {query}")
        print(f"Response: {final_state['messages'][-1].content}\n")
    end_time = time.perf_counter()
    print(f"总耗时：{end_time - start_time}秒")
if __name__ == "__main__":
    asyncio.run(main())